from typing import List, Callable, Union

from langgraph.graph import StateGraph
from langgraph.graph.state import CompiledStateGraph
from redis.asyncio import Redis as AsyncRedis

from ..graph.nodes import BaseNodes
from ..llms import LLMConfig, InsCodeModel
from ..config.fuse import get_langfuse_callback
from .base import BaseGraph


class FlowBaseGraph(BaseGraph):
    workflow: StateGraph = None
    graph: CompiledStateGraph = None
    graph_name: str = ""
    nodes: List[Union[BaseNodes, Callable]] = []
    verbose: bool = False

    def __init__(self,
                 graph_name: str,
                 nodes: List[Union[BaseNodes, Callable]],
                 llm_chat: LLMConfig | InsCodeModel = None,
                 redis: AsyncRedis = None,
                 agent_template: str = None,
                 agent_template_language: str = None,
                 verbose: bool = False,
                 **kwargs):
        super().__init__(graph_name=graph_name, llm_chat=llm_chat, redis=redis, agent_template=agent_template,
                         agent_template_language=agent_template_language, verbose=verbose, **kwargs)
        self.nodes = nodes

    def build(self):
        self.add_nodes()
        self.add_edges()
        langfuse_callback = get_langfuse_callback()
        if langfuse_callback:
            self.graph = self.workflow.compile(checkpointer=self.memory).with_config({"callbacks": [langfuse_callback]})
        else:
            self.graph = self.workflow.compile(checkpointer=self.memory)

    def add_nodes(self):
        for node in self.get_nodes():
            self.workflow.add_node(f"{self.graph_name}{node.__name__}", node)

    def add_edges(self, **kwargs):
        func_nodes = self.get_nodes()
        for i, node in enumerate(func_nodes):
            node_name = f"{self.graph_name}{node.__name__}"

            if i == 0:
                self.workflow.add_edge(self.graph_start, node_name)
            else:
                prev_node = func_nodes[i - 1]
                prev_node_name = f"{self.graph_name}{prev_node.__name__}"
                self.workflow.add_edge(prev_node_name, node_name)

            if i == len(func_nodes) - 1:
                self.workflow.add_edge(node_name, self.graph_end)

    def get_nodes(self) -> []:
        nodes = []
        for node in self.nodes:
            if isinstance(node, BaseNodes):
                for node_info in node.get_nodes():
                    nodes.append(node_info["func"])
            elif callable(node):
                nodes.append(node)
        return nodes
